import os
import tensorflow as tf
from tensorflow import keras
from keras import layers, Sequential

if __name__ == '__main__':
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    # 加载预训练网络模型，include_top=False 表示去掉最后一层
    # resnet = keras.applications.ResNet50(weights='imagenet', include_top=False)
    # resnet.summary()

    # 测试网络的输出
    # input = tf.random.normal([4, 224, 224, 3])
    # output = resnet(input)
    # print(output.shape)

    # 新建一个池化层
    # global_average_layer = layers.GlobalAveragePooling2D()

    # 设置输入数据
    # x = tf.random.normal([4, 7, 7, 2048])
    # result = global_average_layer(x)
    # print(x.shape)
    # print(result.shape)

    # 创建一个全连接层
    # fc = layers.Dense(units=100)
    # x = tf.random.normal([4, 2048])
    # out = fc(x)
    # print(out.shape)

    # model = Sequential(
    #     [
    #         # 将 ResNet50 删除最后一个
    #         keras.applications.ResNet50(weights='imagenet', include_top=False),
    #         # 自定义一个全连接层，输出个数为 100
    #         layers.Dense(units=100)
    #     ]
    # )

    resnet50 = keras.applications.ResNet50(weights='imagenet', include_top=False)
    global_average_layer = layers.GlobalAveragePooling2D()
    fc = layers.Dense(units=100)

    # 冻结 ResNet50 现有的参数，后续只训练现有的 fc 参数
    resnet50.trainable = False

    model = Sequential(
        [resnet50, global_average_layer, fc]
    )

    model.summary()
